from sqlalchemy import create_engine, Column, Integer, String, Text, TIMESTAMP, Boolean, ForeignKey, DateTime
from sqlalchemy.orm import declarative_base, Session, relationship
from pgvector.sqlalchemy import Vector
from datetime import datetime

Base = declarative_base()

class User(Base):
    __tablename__ = 'users'
    
    id = Column(Integer, primary_key=True)
    name = Column(String(100), nullable=False)
    email = Column(String(255), unique=True, nullable=False)
    password_hash = Column(String(255), nullable=False)
    avatar_url = Column(String(500))
    created_at = Column(TIMESTAMP, default=datetime.now)
    updated_at = Column(TIMESTAMP, default=datetime.now, onupdate=datetime.now)
    
    # 关系
    knowledge_bases = relationship("KnowledgeBase", back_populates="owner")
    documents = relationship("Document", back_populates="author")

class KnowledgeBase(Base):
    __tablename__ = 'knowledge_bases'
    
    id = Column(String(50), primary_key=True)  # 使用字符串ID，如 'kb1', 'kb2'
    name = Column(String(255), nullable=False)
    description = Column(Text)
    is_public = Column(Boolean, default=False)
    doc_count = Column(Integer, default=0)
    owner_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    created_at = Column(TIMESTAMP, default=datetime.now)
    updated_at = Column(TIMESTAMP, default=datetime.now, onupdate=datetime.now)
    
    # 关系
    owner = relationship("User", back_populates="knowledge_bases")
    documents = relationship("Document", back_populates="knowledge_base")

class Document(Base):
    __tablename__ = 'documents'
    
    id = Column(String(50), primary_key=True)  # 使用字符串ID，如 'doc1', 'doc2'
    title = Column(String(500), nullable=False)
    type = Column(String(50), nullable=False)  # guide, tutorial, reference, faq, etc.
    status = Column(String(20), default='draft')  # draft, published, archived
    content = Column(Text)
    file_path = Column(String(500))
    file_type = Column(String(20))  # PDF, DOCX, TXT, etc.
    file_size = Column(String(20))
    page_count = Column(Integer)
    knowledge_base_id = Column(String(50), ForeignKey('knowledge_bases.id'), nullable=False)
    author_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    created_at = Column(TIMESTAMP, default=datetime.now)
    updated_at = Column(TIMESTAMP, default=datetime.now, onupdate=datetime.now)
    
    # 关系
    knowledge_base = relationship("KnowledgeBase", back_populates="documents")
    author = relationship("User", back_populates="documents")
    slices = relationship("DocumentSlice", back_populates="document")

class DocumentSlice(Base):
    __tablename__ = 'document_slices'
    
    id = Column(String(50), primary_key=True)  # 使用字符串ID，如 'slice1', 'slice2'
    document_id = Column(String(50), ForeignKey('documents.id'), nullable=False)
    question = Column(Text, nullable=False)
    answer = Column(Text, nullable=False)
    slice_index = Column(Integer, nullable=False)
    created_at = Column(TIMESTAMP, default=datetime.now)
    
    # 关系
    document = relationship("Document", back_populates="slices")

# 保留原有的向量嵌入表
class Embedding(Base):
    __tablename__ = 'embeddings'
    
    id = Column(Integer, primary_key=True)
    name = Column(String(255), nullable=False)
    slice_index = Column(Integer, nullable=False)
    content = Column(Text, nullable=False)
    embedding = Column(Vector(1024), nullable=False)
    created_at = Column(TIMESTAMP, default=datetime.now)
